from typing import Sequence

import numpy as np
from numpy.typing import NDArray
from numpy import (pi, sin, cos, sqrt, conj,
				   real as Re, trace as Tr, abs as Abs)
from scipy.optimize import minimize

__all__ = ['smgoat']

I, sx, sy, sz = np.identity(2, dtype=complex), np.array([[0, 1], [1, 0]], dtype=complex), np.array(
		[[0, -1.j], [1.j, 0.]]), np.diag([1., -1]).astype(complex)

def get_infide_grad_func(infide_type: str, **kwargs):
	match infide_type.lower():
		case 'norm2':
			infide_func = lambda U1, U2: np.clip(Re(Tr((U1 - U2).T.conj() @ (U1 - U2))), 1e-12, np.inf)
			grad_func = lambda U1, U2, dU2_dparam: -2.0 * Re(Tr(U1.T.conj() @ dU2_dparam))
		case 'norm phase2':
			raise NotImplementedError
		case 'corr2' | 'pop2':
			P = np.array([[1.0, 0.0], [0, 0]], dtype=complex) if infide_type.lower() == 'pop2' \
				else kwargs.get('P', I)
			infide_func = lambda U1, U2: np.clip(1.0 - Abs(Tr(P @ U1.T.conj() @ U2) / Tr(P)) ** 2, 1e-12, 1.0)
			grad_func = lambda U1, U2, dU2_dparam: -2.0 / Abs(Tr(P)) ** 2 * Re(
					conj(Tr(P @ U1.T.conj() @ U2)) * Tr(P @ U1.T.conj() @ dU2_dparam))
		case 'pop2 zphase':
			def infide_func(U1: NDArray, U2: NDArray):
				U = U1.T.conj() @ U2
				z_phase = np.angle(U[0, 0]) - np.angle(U[1, 1])
				Uz = cos(z_phase / 2) * I - 1.j * sin(z_phase / 2) * sz
				return np.clip(1.0 - Abs(Tr(U1.T.conj() @ U2 @ Uz) / 2) ** 2, 1e-12, 1.0)

			def grad_func(U1: NDArray, U2: NDArray, dU2_dparam: NDArray):
				U = U2 @ U1.T.conj()
				dU_dparam = dU2_dparam @ U1.T.conj()
				diag0, diag1 = np.diag(U)
				diag0_dparam, diag1_dparam = np.diag(dU_dparam)
				return -1.0 / 4 * 2 * Re(diag0_dparam * diag0.conj() + diag1_dparam * diag1.conj())
		case _:
			raise ValueError
	return infide_func, grad_func

def smgoat(
		thetas0: NDArray,
		phis0: NDArray,
		UT: NDArray,  # target unitary operator
		quad_level: int | Sequence[int],
		max_ple: float = 0.2,
		max_ore: float = 0.2,  # in unit of w_max
		quad_type: str = 'uniform',
		infide_type: str = 'norm2',

		iters: int = 300,
		sparse: bool = False,
		angle_mod: bool = True,
		pulse_len_penalty: float = 0.0,
		var_coeff: float = 0.0,
		verbose: bool = True,
		method: str = 'L-BFGS-B',
		options: dict = None,
		**kwargs
):
	assert (num_pulses := len(thetas0)) == len(phis0)
	"""get infidelity"""
	infide_func, grad_func = get_infide_grad_func(infide_type, **kwargs)
	infide_func = kwargs.get('infide_func', infide_func)
	grad_func = kwargs.get('grad_func', grad_func)
	"""numerical integral setting"""
	if sparse:
		from .smolyak import Smolyak_sparse_grids
		assert isinstance(quad_level, int), "level must be integer when 'sparse=True'"
		nodes, weights = Smolyak_sparse_grids(2, quad_level, quad_type=quad_type)
		del Smolyak_sparse_grids
	else:
		from .smolyak import kron_prod
		assert isinstance(quad_level, Sequence) and len(
				quad_level) == 2, "level must be integer sequences with length 2"
		nodes, weights = kron_prod(quad_level, sym=True, quad_type=quad_type)
		del kron_prod
	nodes *= (max_ple, max_ore)

	# evaluate mean infidelity function
	def eval_infide_and_grad(params: NDArray):
		mean_infide, var_infide = 0., 0.
		mean_grad, var_grad = np.zeros_like(params), np.zeros_like(params)
		for node, weight in zip(nodes, weights):
			if np.isclose(weight, 0.): continue
			ple, ore = node
			ple_a1: float = 1.0 + ple
			f: float = sqrt(ple_a1 ** 2 + ore ** 2)

			tmp_grad = np.zeros_like(params)
			U = I.copy()
			for i in range(num_pulses):
				theta, phi = (params[i], params[num_pulses + i]) if angle_mod else (thetas0[i], params[i])
				h_theta = f * theta / 2.0
				U = (cos(h_theta) * I - 1.j * sin(h_theta) * (
						ple_a1 * cos(phi) * sx + ple_a1 * sin(phi) * sy + ore * sz) / f) @ U
			Ub, Uf = U.copy(), I.copy()
			for i in range(num_pulses):
				theta, phi = (params[i], params[num_pulses + i]) if angle_mod else (thetas0[i], params[i])
				h_theta = f * theta / 2.0
				axis_paulis: NDArray = (ple_a1 * cos(phi) * sx + ple_a1 * sin(phi) * sy + ore * sz) / f

				Ui = cos(h_theta) * I - 1.j * sin(h_theta) * axis_paulis
				dUi_dtheta = f / 2 * (-sin(h_theta) * I - 1.j * cos(h_theta) * axis_paulis)
				dUi_dphi: NDArray = -1.j * sin(h_theta) * (-sin(phi) * sx + cos(phi) * sy) * ple_a1 / f

				Ub = Ub @ Ui.T.conj()
				if angle_mod:
					tmp_grad[i] = grad_func(UT, U, Ub @ dUi_dtheta @ Uf)
					tmp_grad[num_pulses + i] = grad_func(UT, U, Ub @ dUi_dphi @ Uf)
				else: tmp_grad[i] = grad_func(UT, U, Ub @ dUi_dphi @ Uf)
				Uf = Ui @ Uf
			infide = infide_func(UT, U)
			# calculate the mean infidelity and its gradient
			mean_infide += weight * infide
			var_infide += weight * infide ** 2
			mean_grad += weight * tmp_grad
			var_grad += weight * 2 * infide * tmp_grad
		var_infide -= mean_infide ** 2
		var_grad -= 2 * mean_infide * mean_grad

		target, grad = mean_infide + var_coeff * var_infide, mean_grad + var_coeff * var_grad
		if angle_mod:
			target += pulse_len_penalty * sum(params[:num_pulses])
			grad[:num_pulses] += pulse_len_penalty
		return target, grad

	if options is None: options = {}
	result = minimize(
			fun=eval_infide_and_grad,
			x0=np.concatenate([thetas0, phis0]) if angle_mod else phis0,
			method=method,
			jac=True,
			bounds=[(0., 4.0 * pi)] * num_pulses + [(0., 2.0 * pi)] * num_pulses if angle_mod \
				else [(0., 2.0 * pi)] * num_pulses,
			options={
				'maxiter': iters,
				'disp'   : True if verbose else False,
				**options
			},
			**kwargs
	)
	return result
